use num_traits::Num;
use crate::autograd::function::{ForwardArgs, Function};
use crate::autograd::function_ctx::FunctionCtx;
use crate::Tensor;

pub struct MulBackward<T> {
    ctx: FunctionCtx<T>,
}

impl<T> Function<T> for MulBackward<T>
    where
        T: Copy + Num + 'static,
{
    fn new(ctx: FunctionCtx<T>) -> Self {
        Self { ctx }
    }

    fn ctx(&self) -> &FunctionCtx<T> {
        &self.ctx
    }

    fn forward(ctx: &mut FunctionCtx<T>, args: ForwardArgs<T>) -> Tensor<T> {
        if let ForwardArgs::TensorTensor(lhs, rhs) = args {
            let result = lhs.zip(rhs, |a, b| a * b);
            ctx.save_tensors(vec![lhs.clone(), rhs.clone()]);
            return result
        }
        unreachable!()
    }

    fn backward(&self, grad_output: Tensor<T>) -> Vec<Option<Tensor<T>>> {
        let tensors = self.ctx.tensors();
        let grad0 = &tensors[1] * &grad_output;
        let grad1 = &tensors[0] * &grad_output;
        [grad0.into(), grad1.into()].into()
    }
}
